% Thus function solves the RBC model under perfect foresight by solving a
% fixed point problem using the Levenberg Marquardt algorithm for the
% objective function min Q = sqrt(sum(f(x))^2), possibly with upweigthing the 
% residuals as a function of the horizon (if weightOn = 1)

function [Z,f,diffQ,iter,Q,errorMes] = FixedPointSolver_LM(Z0,x_t,y_tN,N,MaxIter,weightOn,setupEPer)

% Unfold setupEPer
ny       = setupEPer.ny;
mx       = setupEPer.mx;
dispOn   = setupEPer.dispOn;

gridPoints = [0.01 0.1 0.3 0.5 0.8 1.0 1.2 1.5 2.0 3.0 5.0];
numGrid    = size(gridPoints,2);

% Setting the dimensions
nzdim2     = ny+mx;
fTmp       = zeros(N*nzdim2,numGrid);
ZvecTmp    = zeros(N*nzdim2,numGrid);
Zvec       = reshape(Z0,N*nzdim2,1);
diffTmp    = zeros(numGrid,1);
diffQ      = 1;
dimf       = N*nzdim2;

if weightOn == 1
    % Upweighting errors in the first 50 time periods
    idx     = reshape(repmat(1:N,nzdim2,1),N*nzdim2,1);
    weights = max(1,1.05.^(N-idx));
else
    weights = ones(N*nzdim2,1);
end
% Evaluting the function at the starting value
f          = DSGEforesight_N(Z0,x_t,y_tN,N,setupEPer);
Q          = sqrt(sum(real(f.*weights).^2+imag(f.*weights).^2,1));

% Starting the optimization
iter            = 0;
errorMes        = 0;
lambda          = setupEPer.lambda0;
updateJ         = 1;
maxNoDecrease   = 10;
countNoDecrease = 0;
scaling         = 1;
while diffQ > setupEPer.tolf && iter < MaxIter && countNoDecrease < maxNoDecrease
    iter = iter + 1;
    if updateJ == 1
        % The gradient
        if setupEPer.JacobianOption == 1
            epsValue   = 1e-8;
            fp         = zeros(dimf,dimf);
            fm         = zeros(dimf,dimf);
            for i=1:dimf
                % Positive step
                ZvecPlus      = Zvec;
                ZvecPlus(i,1) = Zvec(i,1) + epsValue;
                fp(:,i)       = DSGEforesight_N(reshape(ZvecPlus,ny+mx,N),x_t,y_tN,N,setupEPer);
                % Negative step
                ZvecMinus      = Zvec;
                ZvecMinus(i,1) = Zvec(i,1) - epsValue;
                fm(:,i)        = DSGEforesight_N(reshape(ZvecMinus,ny+mx,N),x_t,y_tN,N,setupEPer);
            end
            % Gradient
            J = (fp - fm)/(2*epsValue);
            % For debugging
            %Janal = analJacobian(reshape(Zvec,ny+mx,N),x_t,y_tN,N,setupEPer);;
            %max(max(abs(Janal-J)))
        else
            J = analJacobian(reshape(Zvec,ny+mx,N),x_t,y_tN,N,setupEPer);
        end
        J = real(J).*repmat(weights,1,N*nzdim2);
        AA = J'*J;
    end
    if weightOn == 1
        XX       = (AA + lambda*eye(dimf));
    elseif weightOn == 0
        %XX       = (AA + lambda*diag(diag(AA)));
        XX       = (AA + lambda*eye(dimf));
    end
    condNumber = condest(XX);
    if condNumber > 1D16
        if dispOn ==1
            disp('lambda*diag(I)')
        end
        versionGrid = 2;    % we momentarily use grid search
        while condNumber > 1D16
            lambda     = lambda*2;
            XX         = (AA + lambda*eye(dimf));
            condNumber = condest(XX);
        end
    else
        versionGrid = 1;
    end
    score    = J'*(real(f.*weights));
    delta    = XX\(-score); %J'*(-real(f));
    
    % Evaluating fit at new value of Zvec
    ZvecNew      = Zvec + delta;
    if any(isnan(ZvecNew)) || any(isinf(ZvecNew))
        ZvecNew   = Zvec;
        scaling   = 1;
        QNew      = Q + 1;
    else
        if versionGrid == 1
            fNew     = DSGEforesight_N(reshape(ZvecNew,ny+mx,N),x_t,y_tN,N,setupEPer);
            QNew     = sqrt(sum(real(fNew.*weights).^2+imag(fNew.*weights).^2,1));
            diffQTmp = QNew - Q;
            scaling  = 1;
            idxSearch  = 0;
            while diffQTmp > 0 && idxSearch < 10
                idxSearch = idxSearch + 1;
                scaling   = 0.5^idxSearch;
                ZvecNew   = Zvec + scaling*delta;
                fNew      = DSGEforesight_N(reshape(ZvecNew,ny+mx,N),x_t,y_tN,N,setupEPer);
                QNew      = sqrt(sum(real(fNew.*weights).^2+imag(fNew.*weights).^2,1));
                diffQTmp  = QNew - Q;
            end
        elseif versionGrid == 2
            %The grid search
            for i=1:numGrid
                ZvecTmp(:,i) = Zvec + delta*gridPoints(1,i);
                fTmp(:,i)    = DSGEforesight_N(reshape(ZvecTmp(:,i),nzdim2,N),x_t,y_tN,N,setupEPer);
                diffTmp(i,1) = sqrt(sum(real(fTmp(:,i).*weights).^2+imag(fTmp(:,i).*weights).^2));
                if isnan(diffTmp(i,1)) 
                    diffTmp(i,1) = 1D60;
                end
            end
            [~,idx] = sort(diffTmp);
            fNew    = fTmp(:,idx(1,1));
            QNew    = diffTmp(idx(1,1),1);
            scaling = gridPoints(1,idx(1,1));
            ZvecNew = ZvecTmp(:,idx(1,1));
        end
    end
    diffQ = abs(QNew - Q);
    if dispOn == 1
        disp(['Q                = ',num2str(Q)]);
        disp(['dQ               = ',num2str(QNew-Q)]);
        disp(['lambda           = ',num2str(lambda)]);
        disp(['scaling          = ',num2str(scaling)]);
        disp(['updateJ          = ',num2str(updateJ)]);
        disp(['score length     = ',num2str(sqrt(score'*score))]);
        disp(['Iteration        = ',num2str(iter)]);
        disp(['WNLS             = ',num2str(weightOn)]);
        disp(' ');
    end
    if QNew < Q
        Zvec    = ZvecNew;
        f       = fNew;
        Q       = QNew;
        updateJ = 1;
        countNoDecrease = 0;
    else
        updateJ = 0;
        countNoDecrease = countNoDecrease + 1;
    end
    if scaling < 1
        lambda = lambda*10;  %5
    else
        lambda = lambda/1.5; %1.5
    end
end
if abs(diffQ) > setupEPer.tolf || iter > MaxIter || max(abs(f)) > setupEPer.tolf
    errorMes = 1;
end
Z = reshape(Zvec,nzdim2,N);
% Reporting the objective function without weighting
f          = DSGEforesight_N(reshape(Zvec,nzdim2,N),x_t,y_tN,N,setupEPer);
Q          = sqrt(sum(real(f).^2+imag(f).^2,1));
end
